'''
构建pipeline 的gradio

img1【输入】   
img2【输入】 output_img【输出】
btn_get_res【
    读取img1 =》blured_depth & img2 输出到 output_img 】

使用  
python xxxx.py  cuda_id  port
比如
python gradio_pipeline_for_depth_control.py 4 20021
'''

import argparse

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='示例脚本：接收CUDA设备和端口参数')
    
    # 添加参数
    parser.add_argument(
        '-c',
        '--cuda_id', 
        type=str,
        required=True,  # 必须传入
        help='CUDA设备ID，例如 "0" 或 "0,1"（字符串类型）'
    )
    parser.add_argument(
        '-p',
        '--port',
        type=int,
        default=20025,  # 默认值
        help='端口号（整数类型，默认8000）'
    )
    
    # 解析参数
    args = parser.parse_args()
    return args

args = parse_args()

import os,sys
os.environ['CUDA_VISIBLE_DEVICES']=args.cuda_id

import util_for_huggingface

import gradio as gr
import numpy as np
from PIL import Image,ImageOps
import cv2,requests,io,base64,torch

from diffusers import FluxControlPipeline, FluxPriorReduxPipeline
from util_flux import process_img_1024,vertical_concat_images,horizontal_concat_images

from image_gen_aux import DepthPreprocessor
# from util_mask import load_yolo,get_mask_by_yolo
from demo_rmbg import load_rmbg,get_mask_by_rmbg

from MODEL_CKP import FLUX_DEPTH,FLUX_REDUX,DEPTH_PREDCITION
depth_processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                    FLUX_REDUX, 
                                    torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxControlPipeline.from_pretrained(FLUX_DEPTH, torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights("/data/models/FLUX.1-Turbo-Alpha")
model_rmbg = load_rmbg()


ori_depth = None
now_depth = None
restored_depth = None


from utils.util_for_depthcontrol import get_result,get_result_ori


def _get_result(img1 , img2, 
                alpha = 0.8 ,control_alpha = 0.4,gaussian_score=20):
    return get_result(img1 , img2, 
                depth_processor = depth_processor,
                redux_pipe = pipe_prior_redux, flux_pipe = pipe,
                get_mask = get_mask_by_rmbg, mask_model = model_rmbg,
                alpha = alpha ,control_alpha = control_alpha,gaussian_score=gaussian_score,
                steps=8)
def _get_result_ori(img1 , img2):
    return get_result_ori(img1 , img2,
                depth_processor = depth_processor,
                redux_pipe = pipe_prior_redux, flux_pipe = pipe,
                get_mask = get_mask_by_rmbg, mask_model = model_rmbg,
                steps=8)

with gr.Blocks(title="图像处理Pipeline") as demo:
    gr.Markdown("# 图像处理Pipeline")
    gr.Markdown("上传图片，按步骤处理")
    
    height = 512

    with gr.Row():
        with gr.Column():
            # 第一部分：深度图生成与修复
            img1 = gr.Image(label="输入图像1", 
                            type="pil",height=height)

            btn_get_res_ori = gr.Button("生成原始结果")
            
            output_img_ori = gr.Image(label="输出图像(原始)", type="pil", 
                                  interactive=False,
                                  height=height)
            
        
        with gr.Column():
            # 第二部分：最终图像合成
            img2 = gr.Image(label="输入图像2", type="pil",
                            height=height)

            with gr.Row():
                alpha = gr.Slider(0, 1, value=0.8, label="风格比重")
                control_alpha = gr.Slider(0, 1, value=0.4, label="控制depth比重")
            with gr.Row():
                gaussian_score = gr.Slider(10, 50, value=20,step=5, label="控制高斯模糊的程度")
            

            btn_get_res = gr.Button("生成最终结果(高斯模糊)")
            
            output_img = gr.Image(label="输出图像", type="pil", 
                                  interactive=False,
                                  height=height)
            
    
    # 原始部分处理流程
    btn_get_res_ori.click(
        fn=_get_result_ori,
        inputs=[img1,img2],
        outputs=output_img_ori
    )    
    # 第二部分处理流程
    btn_get_res.click(
        fn=_get_result,
        inputs=[img1,img2,alpha,control_alpha,gaussian_score],
        outputs=output_img
    )
    

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=args.port)